package hbbft

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"
)

// PerformanceMetrics 性能指标收集器
type PerformanceMetrics struct {
	NodeID      uint64
	N           int  // 总节点数
	F           int  // 最大故障节点数
	BatchSize   int  // 每个节点的batch大小
	StartTime   time.Time
	FirstEpochTime time.Time  // 第一个epoch完成的时间
	mu          sync.RWMutex
	
	// 吞吐量统计
	TotalTransactions int
	TotalEpochs       int
	LastEpochTime     time.Time
	
	// 延迟统计
	EpochLatencies    []time.Duration
	ComponentLatencies map[string][]time.Duration
	
	// 每个epoch的交易交付量和延迟（用于TPS计算）
	EpochDeliveries map[uint64]int
	EpochLatencyMap map[uint64]time.Duration  // 每个epoch的延迟
	
	// 组件级延迟
	RBCLatencies    []time.Duration
	BBALatencies    []time.Duration
	ACSLatencies    []time.Duration
	TPKELatencies   []time.Duration
	TBLSLatencies   []time.Duration
	
	// 网络统计
	NetworkStats struct {
		MessagesSent     int
		MessagesReceived int
		BytesSent        int64
		BytesReceived    int64
	}
}

// NewPerformanceMetrics 创建性能指标收集器
func NewPerformanceMetrics(nodeID uint64, n, f, batchSize int) *PerformanceMetrics {
	return &PerformanceMetrics{
		NodeID:            nodeID,
		N:                 n,
		F:                 f,
		BatchSize:         batchSize,
		StartTime:         time.Now(),
		EpochDeliveries:   make(map[uint64]int),
		EpochLatencyMap:   make(map[uint64]time.Duration),
		ComponentLatencies: make(map[string][]time.Duration),
	}
}

// RecordEpochCompletion 记录epoch完成
func (pm *PerformanceMetrics) RecordEpochCompletion(epoch uint64, txCount int, duration time.Duration) {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	
	// 检查该epoch是否已经被记录过，避免重复计数
	if _, exists := pm.EpochDeliveries[epoch]; exists {
		// 该epoch已经记录过，忽略重复记录
		fmt.Printf("[WARNING] Node %d: Duplicate epoch recording detected for epoch %d\n", pm.NodeID, epoch)
		return
	}
	
	pm.TotalEpochs++
	pm.TotalTransactions += txCount
	pm.EpochDeliveries[epoch] = txCount
	pm.EpochLatencyMap[epoch] = duration  // 存储epoch对应的延迟
	pm.EpochLatencies = append(pm.EpochLatencies, duration)
	pm.LastEpochTime = time.Now()
	
	// 记录第一个epoch的时间作为实际测量的起点
	if pm.FirstEpochTime.IsZero() {
		pm.FirstEpochTime = time.Now()
	}
	
}

// RecordComponentLatency 记录组件延迟
func (pm *PerformanceMetrics) RecordComponentLatency(component string, latency time.Duration) {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	
	pm.ComponentLatencies[component] = append(pm.ComponentLatencies[component], latency)
	
	// 同时记录到具体组件
	switch component {
	case "RBC":
		pm.RBCLatencies = append(pm.RBCLatencies, latency)
	case "BBA":
		pm.BBALatencies = append(pm.BBALatencies, latency)
	case "ACS":
		pm.ACSLatencies = append(pm.ACSLatencies, latency)
	case "TPKE":
		pm.TPKELatencies = append(pm.TPKELatencies, latency)
	case "TBLS":
		pm.TBLSLatencies = append(pm.TBLSLatencies, latency)
	}
}

// RecordNetworkMessage 记录网络消息
func (pm *PerformanceMetrics) RecordNetworkMessage(sent bool, bytes int) {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	
	if sent {
		pm.NetworkStats.MessagesSent++
		pm.NetworkStats.BytesSent += int64(bytes)
	} else {
		pm.NetworkStats.MessagesReceived++
		pm.NetworkStats.BytesReceived += int64(bytes)
	}
}

// GetThroughput 获取吞吐量（TPS）
// 根据 Dumbo-NG 论文: TPS = BatchSize * (n-f) / 延迟
// 公式：每个epoch交付 (n-f) 个batch，每个batch包含 BatchSize 笔交易
func (pm *PerformanceMetrics) GetThroughput() float64 {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	
	// 如果没有完成任何epoch，返回0
	if len(pm.EpochLatencies) == 0 || len(pm.EpochDeliveries) == 0 {
		return 0
	}
	
	// 计算平均延迟，跳过异常小的延迟（<50ms可能是计时错误）
	var totalLatency time.Duration
	validCount := 0
	minValidLatency := 50 * time.Millisecond
	
	for _, latency := range pm.EpochLatencyMap {
		if latency >= minValidLatency {
			totalLatency += latency
			validCount++
		}
	}
	
	if validCount == 0 {
		return 0
	}
	
	avgLatency := totalLatency / time.Duration(validCount)
	
	if avgLatency == 0 {
		return 0
	}
	
	// TPS = BatchSize * (n-f) / 平均延迟（秒）
	// 根据 Dumbo-NG 论文，每个epoch交付 (n-f) 个batch，每个batch包含 BatchSize 笔交易
	expectedTxPerEpoch := float64(pm.BatchSize * (pm.N - pm.F))
	return expectedTxPerEpoch / avgLatency.Seconds()
}

// GetAverageLatency 获取平均延迟（跳过异常小的延迟）
func (pm *PerformanceMetrics) GetAverageLatency() time.Duration {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	
	if len(pm.EpochLatencies) == 0 {
		return 0
	}
	
	// 跳过异常小的延迟（<50ms可能是计时错误或异常快的epoch）
	var total time.Duration
	validCount := 0
	minValidLatency := 50 * time.Millisecond
	
	for _, latency := range pm.EpochLatencies {
		if latency >= minValidLatency {
			total += latency
			validCount++
		}
	}
	
	if validCount == 0 {
		return 0
	}
	
	return total / time.Duration(validCount)
}

// GetComponentAverageLatency 获取组件平均延迟
func (pm *PerformanceMetrics) GetComponentAverageLatency(component string) time.Duration {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	
	latencies, exists := pm.ComponentLatencies[component]
	if !exists || len(latencies) == 0 {
		return 0
	}
	
	var total time.Duration
	for _, latency := range latencies {
		total += latency
	}
	return total / time.Duration(len(latencies))
}

// GetStats 获取统计信息
func (pm *PerformanceMetrics) GetStats() map[string]interface{} {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	
	elapsed := time.Since(pm.StartTime)
	
	stats := map[string]interface{}{
		"node_id":            pm.NodeID,
		"elapsed_seconds":    elapsed.Seconds(),
		"total_transactions": pm.TotalTransactions,
		"total_epochs":       pm.TotalEpochs,
		"throughput_tps":     pm.GetThroughput(),
		"avg_epoch_latency":  pm.GetAverageLatency().Milliseconds(),
		"network_stats":      pm.NetworkStats,
		"epoch_deliveries":   pm.EpochDeliveries,
	}
	
	// 添加组件延迟统计
	componentStats := make(map[string]interface{})
	for component, latencies := range pm.ComponentLatencies {
		if len(latencies) > 0 {
			var total time.Duration
			for _, latency := range latencies {
				total += latency
			}
			avgLatency := total / time.Duration(len(latencies))
			componentStats[component] = map[string]interface{}{
				"count":        len(latencies),
				"avg_latency":  avgLatency.Milliseconds(),
				"min_latency":  pm.getMinLatency(latencies).Milliseconds(),
				"max_latency":  pm.getMaxLatency(latencies).Milliseconds(),
			}
		}
	}
	stats["component_latencies"] = componentStats
	
	return stats
}

func (pm *PerformanceMetrics) getMinLatency(latencies []time.Duration) time.Duration {
	if len(latencies) == 0 {
		return 0
	}
	min := latencies[0]
	for _, latency := range latencies[1:] {
		if latency < min {
			min = latency
		}
	}
	return min
}

func (pm *PerformanceMetrics) getMaxLatency(latencies []time.Duration) time.Duration {
	if len(latencies) == 0 {
		return 0
	}
	max := latencies[0]
	for _, latency := range latencies[1:] {
		if latency > max {
			max = latency
		}
	}
	return max
}

// LogStats 输出统计信息到日志
func (pm *PerformanceMetrics) LogStats() {
	stats := pm.GetStats()
	
	// 确保logs目录存在
	logDir := "logs"
	if err := os.MkdirAll(logDir, 0755); err != nil {
		fmt.Printf("Failed to create logs directory: %v\n", err)
		return
	}
	
	// 写入JSON格式的统计信息
	logFile := filepath.Join(logDir, fmt.Sprintf("node_%d_performance.json", pm.NodeID))
	data, err := json.MarshalIndent(stats, "", "  ")
	if err != nil {
		fmt.Printf("Failed to marshal stats: %v\n", err)
		return
	}
	
	if err := os.WriteFile(logFile, data, 0644); err != nil {
		fmt.Printf("Failed to write stats file: %v\n", err)
		return
	}
	
	// 同时输出到控制台
	fmt.Printf("\n=== Performance Stats for Node %d ===\n", pm.NodeID)
	fmt.Printf("Elapsed Time: %.2f seconds\n", stats["elapsed_seconds"])
	fmt.Printf("Total Transactions: %d\n", stats["total_transactions"])
	fmt.Printf("Total Epochs: %d\n", stats["total_epochs"])
	fmt.Printf("Throughput: %.2f TPS\n", stats["throughput_tps"])
	fmt.Printf("Average Epoch Latency: %.2f ms\n", stats["avg_epoch_latency"])
	fmt.Printf("Network - Sent: %d msgs, Received: %d msgs\n", 
		pm.NetworkStats.MessagesSent, pm.NetworkStats.MessagesReceived)
	fmt.Printf("Network - Sent: %d bytes, Received: %d bytes\n", 
		pm.NetworkStats.BytesSent, pm.NetworkStats.BytesReceived)
	
	// 输出组件延迟统计
	if componentStats, ok := stats["component_latencies"].(map[string]interface{}); ok {
		fmt.Printf("\nComponent Latencies:\n")
		for component, compStats := range componentStats {
			if stats, ok := compStats.(map[string]interface{}); ok {
				fmt.Printf("  %s: count=%v, avg=%.2fms, min=%.2fms, max=%.2fms\n",
					component, stats["count"], stats["avg_latency"], 
					stats["min_latency"], stats["max_latency"])
			}
		}
	}
	fmt.Printf("=====================================\n\n")
}


